/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */

package uk.ac.liv;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import javax.xml.parsers.*;
import org.xml.sax.*;
import org.w3c.dom.*;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedHashMap;

/**
 *
 * @author jonesar
 */
public class CrosslinkAnalysis {

    //********** All these should be overwritten by the calling code **************
    private static String inputFile;
    private static String outputMGF;
    private static String outputXLFile;
    private static String inputFastaFile;
    private Double xlMassShift;
    private HashMap<String, Double> varMods = new HashMap();
    private HashMap<String, Double> fixedMods = new HashMap();
    private String xlResidue;                                       //Residue on which crosslinks occur
    private int missedCleaves;
    private Double ppmTol;                                          //Tolerance for the crosslink match    
    //*****************************************************************************
    //********** These are fixed parameters but could be added to the constructor **************
    private Boolean noCrosslinksOnCterminus = true;
    private Boolean noCrosslinksOnModifiedResidues = true;
    //*****************************************************************************
    //private static String inputFastaFile = "resources/myo_tiny.fasta";     
    //private static String inputFile =  "C:/Work/Crosslinking/Waters/BSA_Processed/bsa_processed.xml";
    // private static String inputFastaFile = "C:/Work/Crosslinking/Waters/BSA_Processed/bsa_sequence.fasta";
    //private static String  outputMGF = "C:/Work/Crosslinking/Waters/BSA_Processed/bsa.mgf";    
    //private static String  outputXLFile = "C:/Work/Crosslinking/Waters/BSA_Processed/bsaxl.csv";
    private Document dom;
    private HashMap<Integer, String> lowEnergyPosToHeader = new HashMap();
    private HashMap<Integer, String> highEnergyPosToHeader = new HashMap();
    private HashMap<String, Integer> lowEnergyHeaderToPos = new HashMap();
    private HashMap<String, Integer> highEnergyHeaderToPos = new HashMap();
    private HashMap<String, String[]> lowEnergyData = new HashMap();  //ID --> Split data row map
    private HashMap<String, String[]> highEnergyData = new HashMap();
    private HashMap<String, String[]> parentIDsToFragmentIDs = new HashMap();
    private HashMap<String, Double> potentialMassesFromFasta = new HashMap();
    private HashMap<String, LinkedHashMap<String, Double>> potentialFragmentMassesFromFasta = new HashMap();
    private ArrayList<String> parentIds = new ArrayList();
    private ArrayList<String> parentHeaders = new ArrayList();  //List of column headers found in the data
    private static Double HMASS = 1.007825;          //Mass of H
    private static Double OMASS = 15.994915;        //Mass of O
    private static String path;

    /**
     * @param args the command line arguments
     */
    public static void main(String[] args) {
        // TODO code application logic here

        CrosslinkAnalysis cAnalysis = new CrosslinkAnalysis();


    }

    /**
     * Constructor for testing only
     *
     * @param args the command line arguments
     */
    public CrosslinkAnalysis() {
        //runMyoTests();    //run tests
        runDisulfideTests();
    }

    /*
     * Constructor allowing code to be called from elsewhere and correctly parameterized
     */
    public CrosslinkAnalysis(String inFile, String outMGFFile,
                             String outCrossLinkedFile,
                             Double crosslinkerMassShift,
                             String crosslinkedResidue,
                             HashMap<String, Double> variableMods,
                             HashMap<String, Double> fixMods,
                             int missedCleavages, Double ppmTolerance, String fasta) {
        inputFile = inFile;
        outputMGF = outMGFFile;
        outputXLFile = outCrossLinkedFile;
        xlMassShift = crosslinkerMassShift;

        varMods = variableMods;
        fixedMods = fixMods;
        xlResidue = crosslinkedResidue;
        missedCleaves = missedCleavages;
        ppmTol = ppmTolerance;
        inputFastaFile = fasta;
        init();

    }

    private void init() {

        parseXmlFile(inputFile);
        parseDocument();
        writeMGFFile(outputMGF);
        ProcessFasta pf = new ProcessFasta(inputFastaFile, xlMassShift, xlResidue, missedCleaves, varMods, fixedMods, noCrosslinksOnCterminus, noCrosslinksOnModifiedResidues);
        potentialMassesFromFasta = pf.getAllPeptideMasses();
        potentialFragmentMassesFromFasta = pf.getAllFragmentIonMasses();
        searchForCrosslinkedPeptides(outputXLFile, ppmTol);
    }

    public void runMyoTests() {

        inputFile = "resources/myoC_spectrum.xml";
        outputMGF = "outputFiles/myoC.mgf";
        inputFastaFile = "resources/myoglobin.fasta";

        parseXmlFile(inputFile);
        parseDocument();
        writeMGFFile(outputMGF);

        int missCleave = 2;
        HashMap<String, Double> varMods = new HashMap();
        HashMap<String, Double> fixedMods = new HashMap();
        //For disulphide data: -2.0, C, ss, 1
        varMods.put("K", 156.078644);   //Also test the DSS dead-end as a var mod
        fixedMods.put("C", 57.021464); //Carbamidomethylation

        Double xlMassShift = 156.078644;

        ProcessFasta pf = new ProcessFasta(inputFastaFile, xlMassShift, "K", missCleave, varMods, fixedMods, noCrosslinksOnCterminus, noCrosslinksOnModifiedResidues);
        potentialMassesFromFasta = pf.getAllPeptideMasses();

        Double ppmTol = 5.0;
        outputXLFile = "outputFiles/myo_5.csv";
        searchForCrosslinkedPeptides(outputXLFile, ppmTol);

        ppmTol = 10.0;
        outputXLFile = "outputFiles/myo_10.csv";
        searchForCrosslinkedPeptides(outputXLFile, ppmTol);

        ppmTol = 20.0;
        outputXLFile = "outputFiles/myo_20.csv";
        searchForCrosslinkedPeptides(outputXLFile, ppmTol);

        ppmTol = 50.0;
        outputXLFile = "outputFiles/myo_50.csv";
        searchForCrosslinkedPeptides(outputXLFile, ppmTol);
    }

    public void runDisulfideTests() {

        File directory = new File("");
        path = directory.getAbsolutePath();
        inputFile = "BSA_Processed/bsa_processed.xml";
        outputMGF = path + "/src/main/resources/outputFiles/BSA_Processed/bsa.mgf";
        inputFastaFile = path + "/src/main/resources/BSA_Processed/bsa_sequence.fasta";

        parseXmlFile(inputFile);
        parseDocument();
        writeMGFFile(outputMGF);

        Boolean noCrosslinksOnCterminus = true;
        Boolean noCrosslinksOnModifiedResidues = true;
        int missCleave = 2;
        HashMap<String, Double> varMods = new HashMap();
        HashMap<String, Double> fixedMods = new HashMap();
        //For disulphide data: -2.0, C, ss, 1
        //varMods.put("K", 156.078644);   //Also test the DSS dead-end as a var mod
        //fixedMods.put("C", 57.021464 ); //Carbamidomethylation

        Double xlMassShift = -2.0;
        ProcessFasta pf = new ProcessFasta(inputFastaFile, xlMassShift, "C", missCleave, varMods, fixedMods, noCrosslinksOnCterminus, noCrosslinksOnModifiedResidues);
        potentialMassesFromFasta = pf.getAllPeptideMasses();
        potentialFragmentMassesFromFasta = pf.getAllFragmentIonMasses();

        Double ppmTol = 2.0;
        outputXLFile = path + "/src/main/resources/outputFiles/BSA_Processed/bsa_2.csv";
        searchForCrosslinkedPeptides(outputXLFile, ppmTol);

        ppmTol = 5.0;
        outputXLFile = path + "/src/main/resources/outputFiles/BSA_Processed/bsa_5.csv";
        searchForCrosslinkedPeptides(outputXLFile, ppmTol);

        ppmTol = 10.0;
        outputXLFile = path + "/src/main/resources/outputFiles/BSA_Processed/bsa_10.csv";
        searchForCrosslinkedPeptides(outputXLFile, ppmTol);


        ppmTol = 20.0;
        outputXLFile = path + "/src/main/resources/outputFiles/BSA_Processed/bsa_20.csv";
        searchForCrosslinkedPeptides(outputXLFile, ppmTol);


    }

    private void parseXmlFile(String file) {
        //get the factory
        System.out.println("Parsing XML from file: " + file);
        DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();

        try {

            //Using factory get an instance of document builder
            DocumentBuilder db = dbf.newDocumentBuilder();

            //parse using builder to get DOM representation of the XML file
            dom = db.parse(file);




        }
        catch (ParserConfigurationException pce) {
            pce.printStackTrace();
        }
        catch (SAXException se) {
            se.printStackTrace();
        }
        catch (IOException ioe) {
            ioe.printStackTrace();
        }
    }

    private void parseDocument() {

        //System.out.println("Parsing doc");

        //get the root element
        Element docEle = dom.getDocumentElement();

        NodeList nl = docEle.getElementsByTagName("FORMAT");

        if (nl != null && nl.getLength() > 0) {
            for (int i = 0; i < nl.getLength(); i++) {

                //get the employee element
                Element el = (Element) nl.item(i);

                int fragLevel = Integer.parseInt(el.getAttribute("FRAGMENTATION_LEVEL"));
                if (fragLevel == 0) {
                    processDataHeaders(el, fragLevel);
                }
                else if (fragLevel == 1) {
                    processDataHeaders(el, fragLevel);
                }
                else {
                    System.out.println("Error fragmentation level not recognized: " + fragLevel);
                }
            }
        }

        //get a nodelist of elements
        nl = docEle.getElementsByTagName("DATA");
        if (nl != null && nl.getLength() > 0) {
            Element el = (Element) nl.item(0);
            String data = el.getFirstChild().getNodeValue();
            processSpectralData(data, 0);
        }

        nl = docEle.getElementsByTagName("HE_DATA");
        if (nl != null && nl.getLength() > 0) {
            Element el = (Element) nl.item(0);
            String data = el.getFirstChild().getNodeValue();
            processSpectralData(data, 1);
        }

        nl = docEle.getElementsByTagName("PRECURSOR_BIN");
        joinParentsToFragments(nl);

    }

    private void processDataHeaders(Element formatElement, int fragLevel) {

        /*
         * <FIELD POSITION="1" NAME="Mass" />
         * <FIELD POSITION="2" NAME="Intensity" />
         */
        NodeList nl = formatElement.getElementsByTagName("FIELD");

        if (nl != null && nl.getLength() > 0) {
            for (int i = 0; i < nl.getLength(); i++) {
                Element el = (Element) nl.item(i);

                int headerPos = Integer.parseInt(el.getAttribute("POSITION")) - 1;        //Correct by 1 to start numbering at zero
                String headerName = el.getAttribute("NAME");

                //System.out.println(headerPos + "=>" + headerName);

                if (fragLevel == 0) {
                    parentHeaders.add(headerName);
                    lowEnergyPosToHeader.put(headerPos, headerName);
                    lowEnergyHeaderToPos.put(headerName, headerPos);
                }
                else if (fragLevel == 1) {
                    highEnergyPosToHeader.put(headerPos, headerName);
                    highEnergyHeaderToPos.put(headerName, headerPos);
                }
                else {
                    System.out.println("Error frag level not recognized... :" + fragLevel);
                }
            }
        }

    }

    private void processSpectralData(String data, int fragLevel) {

        String[] peakDataPoints = data.split("\n");

        HashMap<Integer, String> headers;

        if (fragLevel == 0) {
            headers = lowEnergyPosToHeader;
        }
        else {
            headers = highEnergyPosToHeader;
        }

        for (String peakRow : peakDataPoints) {
            peakRow = peakRow.trim();     //Some whitespace on the beginning of the row

            if (peakRow != null && !peakRow.equals("")) {
                //System.out.println(peakRow);
                String[] peakCell = peakRow.split(" ");

                if (fragLevel == 0) {
                    int idPos = lowEnergyHeaderToPos.get("LE_ID");
                    lowEnergyData.put(peakCell[idPos], peakCell);
                }
                else if (fragLevel == 1) {
                    int idPos = highEnergyHeaderToPos.get("HE_ID");
                    highEnergyData.put(peakCell[idPos], peakCell);
                }
                else {
                    System.out.println("Error frag level not recognized: " + fragLevel);
                }
            }
        }
    }

    private void joinParentsToFragments(NodeList nl) {

        //parentIDsToFragmentIDs
        //<PRECURSOR_BIN LE_ID="2" HE_ID="105,106,107" />
        if (nl != null && nl.getLength() > 0) {
            for (int i = 0; i < nl.getLength(); i++) {
                Element el = (Element) nl.item(i);

                String parentID = el.getAttribute("LE_ID");
                String[] fragmentIDs = el.getAttribute("HE_ID").split(",");

                parentIDsToFragmentIDs.put(parentID, fragmentIDs);
                parentIds.add(parentID);                                //To keep ordering of parents by position in the file (i.e. ID)
                //System.out.println(parentID + "==>" + Arrays.toString(fragmentIDs));
            }
        }
    }

    private void writeMGFFile(String outFile) {
        try {
            // Create file 

            FileWriter fstream = new FileWriter(outFile);
            BufferedWriter out = new BufferedWriter(fstream);

            for (String parentId : parentIds) {
                String[] parentData = lowEnergyData.get(parentId);
                String[] fragmentIDs = parentIDsToFragmentIDs.get(parentId);
                //System.out.println("Parent: " + parentId + "data:" + Arrays.toString(parentData));

                String mass = parentData[lowEnergyHeaderToPos.get("Mass")];
                int charge = Integer.parseInt(parentData[lowEnergyHeaderToPos.get("Z")]);
                Double mz = (Double.parseDouble(mass) - HMASS + (charge * HMASS)) / charge;             //They give MH, we need to remove one hydrogen, then add charge * hydrogen, then divide by charge
                //Double mz = (Double.parseDouble(mass))/charge; 
                //688.4114	1374.8082	1377.8344	-3.0262	0	94	3.2e-006	+1	U	K.HGTVVLTALGGILK.K
                //Observed	Mr(expt)	Mr(calc)
                String intensity = parentData[lowEnergyHeaderToPos.get("Intensity")];
                String rt = parentData[lowEnergyHeaderToPos.get("RT")];


                out.write("BEGIN IONS\n");
                out.write("TITLE=CL_" + mz + "_" + rt + "\n");
                out.write("CHARGE=" + charge + "\n");
                out.write("RTINSECONDS=" + rt + "\n");
                out.write("PEPMASS=" + mz + "\n");            //Difficult to work out whether mass in spectrum is the same as intended for MGF
                //In Mascot generic format, the precursor peptide mass is an observed m/z value, from which Mr or MHnn+ is calculated using the prevailing charge state.

                for (String fragmentID : fragmentIDs) {

                    String[] fragmentData = highEnergyData.get(fragmentID);
                    String fragmentMass = fragmentData[highEnergyHeaderToPos.get("Mass")];
                    int fragmentCharge = Integer.parseInt(fragmentData[highEnergyHeaderToPos.get("Z")]);
                    Double fragmentMZ = (Double.parseDouble(fragmentMass) - HMASS + (fragmentCharge * HMASS)) / fragmentCharge;
                    String fragmentIntensity = fragmentData[highEnergyHeaderToPos.get("Intensity")];
                    out.write(fragmentMZ + " " + fragmentIntensity + "\n");
                }
                out.write("END IONS\n\n");
            }
            out.close();
            System.out.println("MGF Output written to " + outFile);
        }
        catch (Exception e) {//Catch exception if any
            System.err.println("Error: " + e.getMessage());
            e.printStackTrace();
        }

    }

    /*
     * Basic search based on MS1 mass
     *
     */
    private void searchForCrosslinkedPeptides(String outFile,
                                              Double toleranceInPPM) {

        try {
            // Create file 
            FileWriter fstream = new FileWriter(outFile);
            BufferedWriter out = new BufferedWriter(fstream);

            // Create detail fragment ions file
            FileWriter fstreamIon = new FileWriter(outFile.replace(".csv", "_ion.csv"));
            BufferedWriter outIon = new BufferedWriter(fstreamIon);

            // Output the header of peptide list
            out.write("Peptide,Theo_MH,Delta,PPM Delta,No. Fragment Ions," + parentHeaders.toString().replace("[", "").replace("]", "") + "\n");

            // Output the header of fragment list
            outIon.write("Peptide,Peptide Mass,Fragment,Fragment Mass,Fragment Theo Mass,Fragment Charge" + "\n");

            for (String parentId : parentIds) {

                String[] parentData = lowEnergyData.get(parentId);
                String[] fragmentIDs = parentIDsToFragmentIDs.get(parentId);

                Double pepMass = Double.parseDouble(parentData[lowEnergyHeaderToPos.get("Mass")]) - HMASS;      //The spectrum contains MH - remove one H, to get pep Mass   

                for (String theoPep : potentialMassesFromFasta.keySet()) {
                    Double theoMass = potentialMassesFromFasta.get(theoPep);
                    Double errorTolInDa = (pepMass) * (toleranceInPPM / 1000000.0);

                    //set the counter for fragment ions
                    int ionCounter = 0;
                    if (pepMass > theoMass - errorTolInDa && pepMass < theoMass + errorTolInDa) {
                        Double delta = (theoMass + HMASS) - Double.parseDouble(parentData[0]);
                        Double percentDelta = delta * 1000000 / theoMass;


                        //for every fragment match we can find                   
                        //out.write("the details of the match\n");
                        for (String fragID : fragmentIDs) {
                            String[] fragData = highEnergyData.get(fragID);
                            //for all fragment masses that can be generated for theoPep (peptide_modstring)
                            //print out if it matches within the given tolerance window
                            Double fragMass = Double.parseDouble(fragData[highEnergyHeaderToPos.get("Mass")]) - HMASS; //remove HMASS assuming ion is +1 charged
                            int fragCharge = Integer.parseInt(fragData[highEnergyHeaderToPos.get("Z")]);
                            HashMap<String, Double> fragIonMasses = potentialFragmentMassesFromFasta.get(theoPep);
                            //HashMap<String, Double> fragIonMasses = potentialFragmentMassesFromFasta.get(theoPep,fragCharge);
                            if (fragIonMasses != null) {
                                for (String fragIonPlusModString : fragIonMasses.keySet()) {
                                    Double theoFragMass = fragIonMasses.get(fragIonPlusModString);
                                    Double fragErrorTolInDa = (fragMass) * (toleranceInPPM / 1000000.0);
                                    if (fragMass > theoFragMass - fragErrorTolInDa && fragMass < theoFragMass + fragErrorTolInDa) {
                                        outIon.write(theoPep + "," + (pepMass + HMASS) + "," + fragIonPlusModString + "," + (fragMass + HMASS.doubleValue()) + "," + (theoFragMass + HMASS.doubleValue()) + ","
                                                + fragCharge + "\n");
                                        //System.out.println("," + fragIonPlusModString + "," + theoFragMass + "," + "\n");
                                        ionCounter++;
                                    }
                                }
                            }
                        }
                        out.write(theoPep + "," + (theoMass + HMASS.doubleValue()) + "," + delta + "," + percentDelta + "," + ionCounter + "," + Arrays.toString(parentData).replace("[", "").replace("]", "") + "\n");
                    }
                }
            }
            out.close();
            outIon.close();
        }
        catch (IOException e) {
            e.printStackTrace();
        }

        System.out.println("Crosslinked output written to " + outFile);

    }

    /**
     * I take a xml element and the tag name, look for the tag and get
     * the text content
     * i.e for <employee><name>John</name></employee> xml snippet if
     * the Element points to employee node and tagName is 'name' I will return John
     */
    private String getTextValue(Element ele, String tagName) {
        String textVal = null;
        NodeList nl = ele.getElementsByTagName(tagName);
        if (nl != null && nl.getLength() > 0) {
            Element el = (Element) nl.item(0);
            textVal = el.getFirstChild().getNodeValue();
        }

        return textVal;
    }

    /**
     * Calls getTextValue and returns a int value
     */
    private int getIntValue(Element ele, String tagName) {
        //in production application you would catch the exception
        return Integer.parseInt(getTextValue(ele, tagName));
    }

}
